# Copyright © 2022 BAAI. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License")

import os
import time
import shlex
import socket
import argparse
import subprocess
from collections import OrderedDict

def run_local_cmd(command, check=True, preview=False, blocking=True, verbose=True):
    ''' Runs command via subprocess in the local machine.

    Args:
        command: command to run on the host
        check: (optional) whether to check the return code of the command
        verbose: (optional) whether to print the command

    Returns:
        result of the command, return code
    '''

    if verbose:
        print("Running command: {} on localhost".format(command), flush=True)

    if not preview:
        if not blocking:
            process = subprocess.Popen(command,
                                       shell=True,
                                       stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE)
            return process, command 
        else:
            result = subprocess.run(command,
                                    shell=True,
                                    capture_output=True,
                                    check=check)
            if verbose:
                if result.returncode != 0:
                    print("Command {} failed on localhost".format(command), flush=True)
                else:
                    print("Command {} successfully executed on localhost".format(command), flush=True)
            return result.stdout.decode("utf-8"), result.returncode
    else:
        return None, None


def run_remote_cmd(host, command, user=None, password=None, blocking=True, preview=False, verbose=True):
    ''' Runs ssh command via subprocess.

    Args:
        host: target host to send the command to
        command: command to run on the host
        user: (optional) user to use to login to host
        password: (optional) password for ssh login
        check: (optional) whether to check the return code of the command
        verbose: (optional) whether to print the command

    Returns:
        result of the command, return code
    '''

    where = "%s" % host if user is None else "%s@%s" % (user, host)
    if password is not None:
        prefix = "sshpass -p {} ".format(password)
    else:
        prefix = ""
    ssh_cmd = "ssh {} '{}'".format(where, command)
    if verbose:
        print("Running command: {} on host {}".format(
            ssh_cmd, host), flush=True)

    if not preview:
        if not blocking:
            process = subprocess.Popen(prefix + ssh_cmd,
                                      shell=True,
                                      stdout=subprocess.PIPE,
                                      stderr=subprocess.PIPE)
            return process, ssh_cmd
        else:
            result = subprocess.run(prefix + ssh_cmd,
                                    shell=True,
                                    capture_output=True,
                                    check=True)
            if verbose:
                if result.returncode != 0:
                    print("Command {} failed on host {}".format(ssh_cmd, host), flush=True)
                else:
                    print("Command {} successfully executed on host {}".format(ssh_cmd,host), flush=True)
            return result.stdout.decode("utf-8"), result.returncode
    else:
        return None, None


def get_ip():
    """Get the IP address of the current host."""
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.settimeout(0)
    try:
        # doesn't even have to be reachable
        s.connect(('10.254.254.254', 1))
        IP = s.getsockname()[0]
    except Exception:
        IP = '127.0.0.1'
    finally:
        s.close()
    return IP

def get_args():
    """Parse commandline."""

    parser = argparse.ArgumentParser(description='Execute a command on remote hosts')
    parser.add_argument('--hostfile', type=str, required=True,
                        help='Path to the hostfile')
    parser.add_argument('--password', type=str, default=None,
                        help='password for ssh login')
    parser.add_argument('--preview', action='store_true',
                        help='Preview the checkpoints to be pulled')
    parser.add_argument('--non-blocking', default=False, action='store_true',
                        help='Execute the command in non-blocking mode')
    parser.add_argument('command', metavar='cmd', type=str,
                        help='command to run on the remote hosts')
    args = parser.parse_args()

    return args
                        

if __name__ == '__main__':
    args = get_args()
    hostfile = args.hostfile
    password = args.password
    preview = args.preview
    command = args.command
    cur_work_dir = os.getcwd()
    cur_host = get_ip()

    # Parse the host file and get host list
    hostlist =[]
    with open(hostfile, 'r') as f:
        for line in f:
            host_and_slot = line.strip().split()
            if len(host_and_slot) > 0:
                host = line.strip().split()[0]
                hostlist.append(host)

    if not args.non_blocking:
        # Run the command on each host
        for host in hostlist:
            if cur_host != host:
                new_command = "cd {} && {}".format(cur_work_dir, command)
                result = run_remote_cmd(
                    host, new_command, password=password, preview=preview)
            else:
                result = run_local_cmd(command, preview=preview)
            if result[1] == 0:
                print(result[0], flush=True)
    else:
        # Run the command on each host
        processes = []
        cmds = []
        for host in hostlist:
            if cur_host != host:
                new_command = "cd {} && {}".format(cur_work_dir, command)
                process, cmd = run_remote_cmd(
                    host, new_command, password=password, preview=preview, blocking=False)
                processes.append(process)
                cmds.append(cmd)
            else:
                process, cmd = run_local_cmd(command, preview=preview, blocking=False)
                processes.append(process)
                cmds.append(cmd)

        visited = {process: False for process in processes} 
        while not all(visited.values()):
            for i, process in enumerate(processes):
                if process is None:
                    visited[process] = True
                    continue
                if process.poll() is not None and visited[process] is False:
                    visited[process] = True
                    _, error = process.communicate()
                    if error:
                        print("Error {} when executing command {}".format(error, cmds[i]), flush=True)
            time.sleep(2)
        print("All commands finished", flush=True)
    